#include <windows.h>
#include <string.h>
#include <stdlib.h>

#include "lib.h"
#include "patch.h"
#include "client.h"
#include "file.h"
#include "patchsrv.h"
#include "pkglist.h"
#include "pkgdesc.h"
#include "serializer.h"
#include "serverfile.h"

#include "rsdloadpkg.h"

#define USE_COMPRESSION

// RSync a requested package (send-only protocol)
//
// Request:
//
// <msg len><F><# blocks><namelen><filename>
//		<21><B><weakchecksum><strongchecksum>
//			.
//			.
//		<3><B>
//
// Responses to this protocol request are:
//
// <3><E><initial error>
//
//	or
//
// <5><S><4byte file size>
//		<msgsize><B><N,C,S><blockid><data block>
//			.
//			.
//		<msgsize><B><N,C,S><blockid><data block>
//			or
//		<3><E><error after initial success>
//
// where <N,C,S> denote if the packet is normal, compressed, or split 
//  and compressed. (compression may cause a packet to become larger 
//  than our out buffer)
//
// Descriptions:
//
// <initial error>
//		EM_ServerDownForMaintenance:  generated in DetermineRequest
//		EM_UnknownPackage:			  requested pkg doesn't exist in list
//		EM_RDSL_OutOfMemory:		  couldn't allocate needed memory
//
//	<error after initial success>
//		EM_RDSL_IOError:			  client sent unexpected message or 
//									  msg corrupt
//		EM_RDSL_FileError:			  unzipped file not available, couldn't
//									  create a port user for the file
//		EM_RDSL_DigestError:		  error during digesting of message
//		EM_RDSL_OutOfMemory:		  couldn't allocate needed memory
//
// No other errors are possible.
//

enum DLPkgStates { DONE, RECV_CHECKSUMS, SEND_CHECKSUM_ACK, 
	SEND_RSYNCBLOCK, CALC_RSYNCBLOCK, SEND_LAST_RSYNCBLOCK};

RSDownloadPackage::RSDownloadPackage(Client* client)
	: Protocol(client)
{
	m_compressor = LZHL_CHANDLE_NULL;
	m_compressed_buffer = (uint8 *)0;
	m_cbqueued = 0;

	// create a compression buffer that can hold the worst case
	//  scenario or our compression algorithm...
	// maxbufsize = bufsize + bufsize/2 + 32
	int32 max_compressed_size;
	max_compressed_size = LZHLCompressorCalcMaxBuf(CServerFile::BLOCK_SIZE + PACKETSIZE);

	// if new fails here and m_compressed_buffer == NULL, we will simply
	//  send uncompressed packets to this client...
	m_compressed_buffer = new uint8[max_compressed_size];

	InitState();
}

RSDownloadPackage::~RSDownloadPackage()
{
	delete [] m_compressed_buffer;
	m_compressed_buffer = (uint8 *)0;

	if(LZHL_CHANDLE_NULL != m_compressor)
	{
		LZHLDestroyCompressor(m_compressor);
		m_compressor = LZHL_CHANDLE_NULL;
	}
}

void 
RSDownloadPackage::InitState()
{
	m_rsync_read = 0;
	memset(m_rsyncbuffer, 0, CServerFile::BLOCK_SIZE + PACKETSIZE);

	m_state = DONE;

	m_current_block = 0;
	m_end_of_file = false;

	m_filename = "";

	m_resumesize = 0;

	m_send_queued = false;

	m_serverfile.Done();		//must delete old checksums before creating new file
	m_serverfile.InitState();

	m_serverfile.SetBuffer(m_rsyncbuffer, 
		(CServerFile::BLOCK_SIZE + PACKETSIZE));

	m_cbqueued = 0;

	if(LZHL_CHANDLE_NULL != m_compressor)
	{
		LZHLDestroyCompressor(m_compressor);
		m_compressor = LZHL_CHANDLE_NULL;
	}

#ifdef USE_COMPRESSION
	// no need to create a compressor if there's no compressed buffer or
	//  client... if this fails, we send uncompressed data to the client.
	//  the return value of LZHLCreateCompressor is ignored intentionally
	//  for this reason.
	if(client && m_compressed_buffer)
	{
		// NOTE: the compression tables persist between calls to 
		//  CreateCompressor and DestroyCompressor.  note that we create 
		//  a new one for each client request which means we create
		//  a new table for each file.  b/c of the varied patterns for
		//  the different files we're transmitting (.exe, .rsc, .hlp)
		//  this seemed appropriate.
		m_compressor = LZHLCreateCompressor();
	}
#endif
}

bool8
RSDownloadPackage::recognize(uint8* Msg, uint32 Len)
{
	bool8 success = false;

//	return (bool8)(Msg[0] == 'F' && Len > 1);
	m_multibyte = false;
	success = (bool8)(Msg[0] == 'F' && Len > 1);
	if (!success)
	{
		success = (bool8)(((*(uint16*)&Msg[0]) == *(uint16*)"MF")&& Len >= 2);
		if (success)
			m_multibyte = true;
	}
	return (success);
}

int32
RSDownloadPackage::start(uint8* msg, uint32 len)
{
	int32 result = EM_UnknownPackage;

	// since this protocol can be called multiple times between the 
	//  construction and deletion of the client (and therefore ourselves)
	//  we must init our state with each recognized protocol.
	InitState();

	uint8 *curr = msg;

	uint8 offset = 1;
	if ((*(uint16*)curr == *(uint16*)"MF"))
		offset = 2;

	if(*curr == 'F' || (*(uint16*)curr == *(uint16*)"MF"))
	{
		uint32 namelen = 0;
		uint32 numblocks = 0;
		curr += offset;
		curr += CSerializer::UnWire(curr, &m_resumesize);
		curr += CSerializer::UnWire(curr, &numblocks);

		if(m_serverfile.SetNumBlocks(numblocks))
		{
			if (m_multibyte)
			{
				curr += m_filename.deserialize(curr);
			}
			else
			{
				curr += CSerializer::UnWire(curr, &namelen);

				if (namelen - 1 >= MAX_PATH)
					namelen = MAX_PATH;

				char* name = new char[namelen+1];
				strncpy(name, (char*)curr, namelen);
				name[namelen] = '\0';
				m_filename = name;
				delete name;
			}


			DEBUGMSG(("%S requested.", (const uint16*)m_filename));

			if (m_filename.find(L"%") == 0)
			{
				// a percent at the front of the requested
				//  file means that the first path is the name of
				//  the package the requested file belongs to and should
				//  be replaced with the zipfilename.  the client doesn't
				//  know the zipfilenames...

				// find the pkg name
				TCHAR temp[MAX_PATH+1];
				wcscpy(temp, m_filename);
				TCHAR *c, *rest;

				rest = const_cast<TCHAR*>( wcschr(temp, L'/') );
				if(!rest || (wcschr(temp, L'\\') < rest))
					rest = wcschr(temp, L'\\');

				if(rest)
					*rest = '\0';
				else
					return EM_RSDL_FileError;

				PackageList *list;
				PackageDescription *item;

				list = PatchServer::This->getPackageList();
				if(list)
				{
					int32 count = list->getCount();
					for(int32 i = 0; i < count; i++)
					{
						item = list->getItem(i);
						if(!_wcsicmp(temp+1, item->getName() ))
						{
							// found it!
							const TCHAR *pathname = wcsrchr(item->getZipName(), L'\\');
							if(!pathname || (wcsrchr(item->getZipName(), L'/') > pathname))
							{
								pathname = wcsrchr(item->getZipName(), L'/');
							}
							if(!pathname)
							{
								PatchServer::This->releasePackageList();
								return EM_RSDL_FileError;
							}

							pathname++;

							wcscpy(temp, pathname);
							c = wcsrchr(temp, L'.');
					
							if(c)
							{
								*c = '\0';
							}
							else
							{
								PatchServer::This->releasePackageList();
								return EM_RSDL_FileError;
							}

							rest = const_cast<TCHAR*>( wcschr(m_filename, L'/') );
							if(!rest || (wcschr(m_filename, L'\\') < rest))
								rest = const_cast<TCHAR*>( wcschr(m_filename, L'\\') );

							if(!rest)
							{
								PatchServer::This->releasePackageList();
								return EM_RSDL_FileError;
							}

							wcscpy(temp, rest+1);
							wcscpy((unsigned short*)(const unsigned short*)m_filename, temp);

							client->SetPackage(item);

							break;
						}
					}

					PatchServer::This->releasePackageList();
				}
				else
				{
					return EM_RSDL_FileError;
				}

			}

			result = EM_Success;

			m_state = RECV_CHECKSUMS;
		}
		else
		{
			result = EM_RSDL_OutOfMemory;
		}

	}

	return result;
}


void
RSDownloadPackage::BeginRsync()
{
	TCHAR pathname[MAX_PATH];
    int nStrLen;

	m_rsync_read = 0;
	memset(m_rsyncbuffer, 0, (CServerFile::BLOCK_SIZE + PACKETSIZE));

	// the client knows the package we want...
	if(!client->GetPackage() || !client->GetPackage()->getZipName().getLength())
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_FileError);
		return;
	}		

	wcscpy(pathname, client->GetPackage()->getZipName());

	// the unzipped version is kept in a directory of the same
	//  name as the zip file.
    nStrLen = wcslen(pathname);
    if (nStrLen <= 4 || wcsicmp(pathname+nStrLen-4, L".zip") != 0)
    {
		m_state = DONE;
		client->errmsg(EM_RSDL_FileError);
		return;
    }
	pathname[nStrLen - 4] = '\0';

	// and now we concat the filename that we're interested in...
	wcscat(pathname, L"\\");
	wcscat(pathname, m_filename);

	// and viola!  the uncompressed file we're interested in...
	wcscpy(m_filename.getBuffer(wcslen(pathname)+1), pathname);
	m_filename.releaseBuffer();

	// open the file and find out the size (we don't actually
	//  post a read until we send some info to the client...)
	if(!client->getFile()->open(m_filename))
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_FileError);
		return;
	}

	if(!PatchServer::This->getPortIO()->newPortUser(
		client->getFile()))
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_FileError);
		return;
	}

	// seek to the point the client wishes to start receiving
	if(m_resumesize)
	{
		// begin dishing out bytes and blocks from this position.
		client->getFile()->seek(m_resumesize);
	}

	// now we send an ack to the client that we're beginning to digest
	//  their request...
	uint8* obuf = client->getOBuf();
	uint8* curr = obuf + 1;
	uint32 filesize = client->getFile()->getSize();

	curr += CSerializer::Wire(curr, (uint8)'S');
	curr += CSerializer::Wire(curr, filesize);

    ASSERT( curr - obuf < 256 );
	obuf[0] = curr - obuf;

	m_state = SEND_CHECKSUM_ACK;

	client->setOUsed(curr - obuf);
	client->writeNet();

}

void
RSDownloadPackage::receiveChecksums(uint8 *buffer, uint32 len)
{
	while(len)
	{
		uint32 used 		= 0;
		uint32 total_used 	= 0;
		uint32 msg_size		= 0;

		uint8 *curr;
		curr = buffer;

		msg_size = *curr++;

		if(len < msg_size) 
		{
			break;
		}

		total_used += 1;

		while(total_used < msg_size)
		{
			used = 0;
			if('B' == *curr++)
			{
				if(msg_size == 0x03)
				{
					// this is all of the checksums
					// we're now ready to start reading the file in
					//  and calculating the rsync blocks to send to the
					//  client.
					used += 0x02;
					m_state = CALC_RSYNCBLOCK; 

					// tell the server file we're done receiving the checksums
					//  form the client and we're ready to sort them all and
					//  setup the hash table, etc.

					if(!m_serverfile.DoneChecksums())
					{
						m_state = DONE;
						client->eatIBuf(client->getIUsed());
						client->errmsg(EM_RSDL_DigestError);
						return;
					}

					client->eatIBuf(total_used + used);

					// and now begin reading the file in and digesting it!
					BeginRsync();
					return;
				}
				else
				{
					if(m_serverfile.m_md4_checksums && 
						m_serverfile.m_weak_checksums)
					{
						// the weak and strong checksums follow...

						// WEAK
						curr += CSerializer::UnWire(curr, 
							&m_serverfile.m_weak_checksums[m_current_block]);

						// STRONG
						curr += CSerializer::UnWire(curr, 
							&m_serverfile.m_md4_checksums[m_current_block][0]);
						curr += CSerializer::UnWire(curr, 
							&m_serverfile.m_md4_checksums[m_current_block][1]);
						curr += CSerializer::UnWire(curr, 
							&m_serverfile.m_md4_checksums[m_current_block][2]);
						curr += CSerializer::UnWire(curr, 
							&m_serverfile.m_md4_checksums[m_current_block][3]);
					}
					else
					{
						m_state = DONE;
						client->eatIBuf(client->getIUsed());
						client->errmsg(EM_RSDL_OutOfMemory);	
						return;
					}

					used += 21;
					m_current_block++;

				}
			}
			else
			{
				m_state = DONE;
				client->eatIBuf(client->getIUsed());
				client->errmsg(EM_RSDL_IOError);	
				return;
			}

			total_used += used;
		}

		client->eatIBuf(total_used);

		len 	-= total_used;

	}

	m_state = RECV_CHECKSUMS;
	client->readNet();
}

void
RSDownloadPackage::netIn(int32 amt)
{
	//if (client->bytespersec != -1)
	//	__asm { int 3 };

	if(!amt && m_state != CALC_RSYNCBLOCK)
	{
		// first call
		m_state = RECV_CHECKSUMS;

		client->readNet();

		return;
	}

	uint8 *curr = client->getIBuf();
	switch(m_state)
	{
		// we're still receiving a list of checksums from the client for
		//  this file
		case RECV_CHECKSUMS:
			receiveChecksums(curr, client->getIUsed());
			break;

		case CALC_RSYNCBLOCK:
			if (amt != 0)
			{
				int32 ieat = 0, overhead = 1;
				if (client->getIBuf()[0] == 0)
				{
					overhead = 3;
					ieat = client->msgGet(client->getIBuf()+1, (uint16*)0)+1;
				}
				else
					ieat = client->msgGet(client->getIBuf(), (uint8 *)0);
				
				uint8* msg = client->getIBuf() + overhead;
				
				if (msg[0] == 'S')
				{
					int32 bytespersecond = client->msgGet(&msg[1], (uint16 *)0);
					
					DEBUGMSG(("change download speed: %d", bytespersecond));

					client->setBytesPerSecond(bytespersecond);
					
					client->eatIBuf(amt);
				}
			}
			
			// m_rsync_read is the amount we haven't digested yet.
			client->getFile()->read(m_rsyncbuffer + m_rsync_read, 
				CServerFile::BLOCK_SIZE + PACKETSIZE - m_rsync_read);
			
			break;

		default:
			m_state = DONE;
			client->errmsg(EM_RSDL_IOError);	
			break;
	}
}

int32
RSDownloadPackage::CompressPacket(uint8 *buf, int32 bufsize)
{
	// make sure we have a valid compressor...
	if(LZHL_CHANDLE_NULL != m_compressor && m_compressed_buffer)
	{
		// compress the packet 
		// allocate our compression block
		int32 compressed_size;
		compressed_size = 
			LZHLCompress(m_compressor, m_compressed_buffer, buf, bufsize);

		return compressed_size;
	}

	return -1;
}

bool8 
RSDownloadPackage::SendRSyncBlock()
{
	try
	{
		// now calculate the size of all the raw data and fill in those
		//  things we skipped...
		uint8 *obuf 	= client->getOBuf();
		uint8 *endbuf 	= m_serverfile.GetOBuf();
		int32 bufsize	= (endbuf - obuf) - BLOCK_HEADERSIZE;
		int32 compressed_size;
		int32 blockid = m_serverfile.GetBlockID();

		compressed_size = CompressPacket(obuf + BLOCK_HEADERSIZE, bufsize);
		if(compressed_size >= 0)
		{
			int32 maxsendsize = PACKETSIZE - BLOCK_HEADERSIZE - 1;

			if(compressed_size <= maxsendsize)
			{
				memcpy(obuf + BLOCK_HEADERSIZE, 
					m_compressed_buffer, compressed_size);
			}
			else
			{
				// the block compressed to larger than our send queue...
				//  this rarely happens, so i didn't take a lot of time
				//  to optimize this.  it sends what it can and queues
				//  the rest for another packet.
				// NOTE:
				//  because the size of packets is contained in one block,
				//  the size for ONE packet is limited to the range of [0-255]
				//  so we can only send (PACKETSIZE - 1) bytes on this
				//  first packet...

				memcpy(obuf + BLOCK_HEADERSIZE, 
					m_compressed_buffer, maxsendsize);

				memmove(m_compressed_buffer, 
					m_compressed_buffer + maxsendsize, 
					compressed_size - maxsendsize);

				m_send_queued = true;
				m_cbqueued = compressed_size - maxsendsize;
				blockid = -1;
			}
		}

		if(compressed_size >= 0 && m_send_queued)
		{
			obuf[0] = (uint8)((PACKETSIZE-1));
			obuf[1] = 'B';
			obuf[2] = 'S';
		}
		else if(compressed_size >= 0 && !m_send_queued)
		{
			obuf[0] = compressed_size + BLOCK_HEADERSIZE;
			obuf[1] = 'B';
			obuf[2] = 'C';
		}
		else
		{
			ASSERT(endbuf - obuf < 256);

			obuf[0] = endbuf - obuf;
			obuf[1] = 'B';
			obuf[2] = 'N';
		}

		CSerializer::Wire(&obuf[3], blockid);

		client->setOUsed(obuf[0]);
		client->writeNet();

		return true;
	}
	catch(...)
	{
		return false;
	}
}

int32 
RSDownloadPackage::Digest(uint32 amount)
{
	int32 digested;

	uint8 *obuf = client->getOBuf();
	int32 obufsize = PACKETSIZE;

	obuf = obuf + BLOCK_HEADERSIZE;
	obufsize -= BLOCK_HEADERSIZE;

	digested = m_serverfile.Digest(m_rsync_read, obuf, obufsize);

	return digested;
}

void 
RSDownloadPackage::sendQueued()
{
	int32 blockid;
	uint8 *obuf = client->getOBuf();

	memcpy(obuf + BLOCK_HEADERSIZE, 
		m_compressed_buffer, m_cbqueued);

	m_send_queued = false;
	blockid = m_serverfile.GetBlockID();

	obuf[0] = m_cbqueued + BLOCK_HEADERSIZE;
	obuf[1] = 'B';
	obuf[2] = 'C';

	CSerializer::Wire(&obuf[3], blockid);

	client->setOUsed(obuf[0]);
	client->writeNet();
}

void
RSDownloadPackage::netOut(int32 /*Amt*/)
{
	switch(m_state)
	{
		case SEND_CHECKSUM_ACK:
			// switch states to start calculating the rsync block
			m_state = CALC_RSYNCBLOCK;

			// and ask for our first block read...
			client->getFile()->read(m_rsyncbuffer, 
				CServerFile::BLOCK_SIZE + PACKETSIZE);

			break;

		case SEND_RSYNCBLOCK:
			if(m_send_queued)
			{
				sendQueued();
				return;
			}
			else
			{
				// we just sent the block out to the client, so now
				//  we read in some more...

				m_state = CALC_RSYNCBLOCK;

				if (client->isDatatoRead(6, 3, 'S'))
					client->readNet();
				else
					netIn(0);
			}
			break;

		case SEND_LAST_RSYNCBLOCK:
			if(m_send_queued)
			{
				sendQueued();
				return;
			}
			else
			{
				// if there is still undigested data, continue sending
				if(m_rsync_read)
				{
					int32 digested;
					// and finish digesting the rest of the file which presumably 
					//  is less than an entire block.
					digested = Digest(m_rsync_read);

					if(digested < 0)
					{
						m_state = DONE;
						client->errmsg(EM_RSDL_DigestError);
						return;
					}
					else
					{
						m_rsync_read -= digested;
						memmove(m_rsyncbuffer, m_rsyncbuffer + digested, 
							(CServerFile::BLOCK_SIZE + PACKETSIZE) - digested);

						// and set the rest of the buffer to zeroes
						memset(
							m_rsyncbuffer +
							((CServerFile::BLOCK_SIZE + PACKETSIZE) - digested),
							0, digested);
					}
					m_state = SEND_LAST_RSYNCBLOCK;

					if(!SendRSyncBlock())
					{
						m_state = DONE;
						client->errmsg(EM_RSDL_IOError);
						return;
					}
				}
				else
				{
					// and that's it!
					m_state = DONE;
					client->getFile()->close();
					client->protocolDone();

					DEBUGMSG(("client done."));
				}
			}

			break;

		default:
			m_state = DONE;
			client->errmsg(EM_RSDL_DigestError);
			return;
			break;
	}
}

void
RSDownloadPackage::fileIn(int32 Amt)
{
	if(CALC_RSYNCBLOCK == m_state)
	{
		int32 digested = -1;
		m_rsync_read += Amt;

		// if we have more than a block, then we can digest the block
		//  and send updates to the client.  we can also digest it if
		//  we're at the end of the file even if we don't have a full
		//  block.
		if(m_rsync_read > CServerFile::BLOCK_SIZE)
		{
			// tell the server file to digest what we have so far.
			digested = Digest(m_rsync_read);

			if(digested < 0)
			{
				m_state = DONE;
				client->errmsg(EM_RSDL_DigestError);
				return;
			}
			else
			{
				// and eat the amount digested
				m_rsync_read -= digested;
				memmove(m_rsyncbuffer, m_rsyncbuffer + digested, 
					(CServerFile::BLOCK_SIZE + PACKETSIZE) - digested);

				// and set the rest of the buffer to zeroes
				memset(
					m_rsyncbuffer +
					((CServerFile::BLOCK_SIZE + PACKETSIZE) - digested),
					0, digested);
			}

			m_state = SEND_RSYNCBLOCK;

			// send that baby!	
			if(!SendRSyncBlock())
			{
				m_state = DONE;
				client->errmsg(EM_RSDL_IOError);
				return;
			}

			//  make sure no operations occur after the netout post...
			return;
		}
		else if(!m_end_of_file)
		{
			// we didn't have enough info to digest the block, so just
			//  ask for more more...
			m_state = CALC_RSYNCBLOCK;

			client->getFile()->read(m_rsyncbuffer + m_rsync_read, 
				CServerFile::BLOCK_SIZE + PACKETSIZE - m_rsync_read);
		}
	}
	else
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_FileError);
		return;
	}
}

void
RSDownloadPackage::fileEOF()
{
	// just b/c we've reached the end of the file doesn't mean we're done..
	//  we still have to digest the remaining blocks in our buffer...
	m_end_of_file = true;

	int32 digested;

	// and finish digesting the rest of the file which presumably is less
	//  than an entire block.
	digested = Digest(m_rsync_read);

	if(digested < 0)
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_DigestError);
		return;
	}
	else
	{
		m_rsync_read -= digested;
		memmove(m_rsyncbuffer, m_rsyncbuffer + digested, 
			(CServerFile::BLOCK_SIZE + PACKETSIZE) - digested);

		// and set the rest of the buffer to zeroes
		memset(
			m_rsyncbuffer +
			((CServerFile::BLOCK_SIZE + PACKETSIZE) - digested),
			0, digested);
	}

	m_state = SEND_LAST_RSYNCBLOCK;

	if(!SendRSyncBlock())
	{
		m_state = DONE;
		client->errmsg(EM_RSDL_IOError);
		return;
	}
}

void
RSDownloadPackage::fileErr(uint32 Err)  // kharmon.  04-05-07.  #19307.  Change argument to uint32 so that it overrides base class version.
{
	m_state = DONE;
	client->errmsg(EM_RSDL_FileError);
}



